## Copyright 2002-2004 Andrew Loewenstern, All Rights Reserved
## Copyright 2007 Roee Shlomo, All Rights Reserved
# see LICENSE.txt for license information

from khash import intify
from utility import decodeNodes, decodePeers
from const import *
from errors import *
    
DEBUG = False

class ActionBase:
    """ base class for some long running asynchronous proccesses like finding nodes or values """
    def __init__(self, factory, target, callback):
        self.factory = factory
        self.target = target
        self.num = intify(target)
        self.found = {}
        self.queried = {}
        self.answered = {}
        self.callback = callback
        self.outstanding = 0
        self.finished = False

    def close(self):
        self.finished = True
        self.found = None
        self.queried = None
        self.answered = None
                    
class FindNode(ActionBase):
    """
    Find Node action: nodes lookup
    """
    def goWithNodes(self, nodes):
        """
        This starts the process, our argument is a list of nodes
        """
        for node in nodes:
            if node.id != self.factory.node.id:
                self.found[node.id] = node
                
        self.schedule()
        
    def schedule(self):
        """
        Send messages to new peers, if necessary
        """
        if self.finished:
            return
        l = self.found.values()
        l.sort(key = lambda x: self.num ^ x.num)
        for node in l[:K]:
            if node.id not in self.queried and node.id != self.factory.node.id:
                try:
                    df = node.findNode(self.target, self.factory.node.id)
                except KrpcGenericError:
                    pass
                else:
                    df.addCallback(self.handleGotNodes)
                    df.addErrback(self.defaultGotNodes, node=node)
                    self.outstanding += 1
                    self.queried[node.id] = 1
            if self.outstanding >= CONCURRENT_REQS:
                break
        assert(self.outstanding) >= 0
        if self.outstanding == 0:
            self.factory.rawserver.add_task(self.callback, 0, [l[:K]])
            self.close()
    
    def handleGotNodes(self, dict):
        """
        Called when got respond with a list of nodes
        """
        _krpc_sender = dict['_krpc_sender']
        dict = dict['rsp']
        try:
            l = dict['nodes']
            l = decodeNodes(l)
        except:
            l  = []
##        if DEBUG:
##            print ">>> got nodes %s/%s" % _krpc_sender, len(l)
        sender = self.factory.Node().init(dict["id"], *_krpc_sender)
        self.factory.table.insertNode(sender)
        if self.finished or sender.id in self.answered:
            return
        self.outstanding -= 1
        self.answered[sender.id] = 1
        for node in l:
            n = self.factory.Node().initWithDict(node)
            if n.id not in self.found and n.id not in self.queried:
                self.found[n.id] = n
                self.factory.insertNode(n, contacted = False)
        self.schedule()

    def defaultGotNodes(self, err, node):
        """
        Called when got error
        """
##        if DEBUG:
##            print ">>> find failed %s/%s" % (node.host, node.port), err
        self.factory.table.nodeFailed(node)
        del self.found[node.id] # don't include known stale nodes in the found list
        self.outstanding -= 1
        self.schedule()           
    
class GetValue(FindNode):
    """
    Get Value action: retrieve values from global table
    """
    def __init__(self, factory, target, callback, find="findValue", donecallback = None):
        FindNode.__init__(self, factory, target, callback)
        self.findValue = find
        self.donecallback = donecallback
            
    def goWithNodes(self, nodes, found=[]):
        """
        This starts the process, our argument is a list of nodes
        """
        self.results = {}
        for n in found:
            self.results[n] = 1
            
        for node in nodes:
            if node.id != self.factory.node.id:
                self.found[node.id] = node
            
        self.schedule()

    def handleGotNodes(self, dict):
        """
        Called when got respond with a nodes or value
        """
        _krpc_sender = dict['_krpc_sender']
        dict = dict['rsp']
        sender = self.factory.Node().init(dict["id"], *_krpc_sender)
        self.factory.table.insertNode(sender)
        if self.finished or sender.id in self.answered:
            return

        self.outstanding -= 1
        self.answered[sender.id] = sender
        # go through nodes
        # if we have any closer than what we already got, query them
        if 'nodes' in dict:
            try:
                nodes = decodeNodes(dict['nodes'])
            except:
                del self.answered[sender.id]
            else:
                if DEBUG:
                    print ">>> get nodes success %s/%s, %d" % (sender.host, sender.port, len(nodes))
                for node in nodes:
                    try:
                        n = self.factory.Node().initWithDict(node)
                    except KrpcGenericError:
                        pass
                    else:
                        if n.id not in self.found and n.id not in self.queried:
                            self.found[n.id] = n
        elif 'values' in dict:
            def x(y, z=self.results):
                if y not in z:
                    z[y] = 1
                    return y
                else:
                    return None
            values = decodePeers(dict['values'])
            z = len(values)
            v = filter(None, map(x, values))
            if DEBUG:
                print ">>> get values success %s/%s, %d" % (sender.host, sender.port, len(v))
            if(len(v)):
                self.factory.rawserver.add_task(self.callback, 0, [v])
        self.schedule()
        
    def defaultGotNodes(self, err, node):
        """
        Called when got error
        """
        from traceback import print_exc
        if DEBUG:
            print ">>> get failed %s/%s" % (node.host, node.port), err

        self.factory.table.nodeFailed(node)
        del self.found[node.id] # don't include known stale nodes in the found list
        self.outstanding -= 1
        self.schedule()
            
    def schedule(self):
        """
        Send messages to new peers, if necessary
        """
        if self.finished:
            return
        l = self.found.values()
        l.sort(key = lambda x: self.num ^ x.num)
        
        for node in l[:K]:
            if node.id not in self.queried and node.id != self.factory.node.id:
                try:
                    f = getattr(node, self.findValue)
                except AttributeError:
                    if DEBUG:
                        print ">>> findValue %s doesn't have a %s method!" % (node, self.findValue)
                else:
                    try:
                        df = f(self.target, self.factory.node.id)
                        df.addCallback(self.handleGotNodes)
                        df.addErrback(self.defaultGotNodes, node=node)
                        self.outstanding += 1
                        self.queried[node.id] = 1
                        if DEBUG:
                            print ">>> get request %s/%s" % (node.host, node.port)
                    except KrpcGenericError:
                        pass
            if self.outstanding >= CONCURRENT_REQS:
                break
        assert(self.outstanding) >= 0
        if self.outstanding == 0:
            if DEBUG:
                print ">>> get done %s" % len(self.answered)
            ## all done
            if self.donecallback:
                self.factory.rawserver.add_task(self.donecallback, 0, [self.answered.values()])
            else:
                self.factory.rawserver.add_task(self.callback, 0, [[]])
            self.close()

class StoreValue(ActionBase):
    """
    Store Value action: store values in the global table
    """
    def __init__(self, factory, target, value, callback, store="storeValue"):
        ActionBase.__init__(self, factory, target, callback)
        self.value = value
        self.stored = []
        self.store = store
        
    def goWithNodes(self, nodes):
        self.nodes = nodes
        self.nodes.sort(key = lambda x: self.num ^ x.num)
        self.schedule()

    def schedule(self):
        if self.finished:
            return
        num = CONCURRENT_REQS - self.outstanding
        if num > STORE_REDUNDANCY:
            num = STORE_REDUNDANCY
        for i in range(num):
            try:
                node = self.nodes.pop()
            except IndexError:
                if self.outstanding == 0:
                    self.close()                    
                    if self.callback:
                        self.callback(self.stored)
            else:
                if not node.id == self.factory.node.id:
                    self.outstanding += 1
                    try:
                        f = getattr(node, self.store)
                    except AttributeError:
                        if DEBUG:
                            print ">>> %s doesn't have a %s method!" % (node, self.store)
                    else:
                        try:
                            df = f(self.target, self.value, self.factory.node.id)
                        except KrpcGenericError:
                            pass
                        except KrpcProtocolError:
                            self.factory.table.invalidateNode(node)
                        else:
                            df.addCallback(self.storedValue, node=node)
                            df.addErrback(self.storeFailed, node=node)

    def storedValue(self, t, node):
        self.outstanding -= 1
        self.factory.table.insertNode(node)
        if self.finished:
            return
        self.stored.append(t)
        if len(self.stored) >= STORE_REDUNDANCY:
            self.close()
            self.callback(self.stored)
        else:
            if not len(self.stored) + self.outstanding >= STORE_REDUNDANCY:
                self.schedule()
        return t
    
    def storeFailed(self, t, node):
        if DEBUG:
            print ">>> store failed %s/%s" % (node.host, node.port)
        self.factory.table.nodeFailed(node)
        self.outstanding -= 1
        if self.finished:
            return t
        self.schedule()
        return t
